import os
from pathlib import Path
import time
import torch
from PIL import Image
from tqdm import tqdm
from torch.utils.data import Dataset
from torchvision import transforms
from collections import Counter

# 检查关键文件是否存在
def verify_dataset(root="E:/codes/project/datas/VeRi-776"):
    required_files = [
        "name_train.txt", "name_test.txt", "name_query.txt",
        "camera_ID.txt", "image_train", "image_test", "image_query"
    ]
    
    missing = []
    for f in required_files:
        path = Path(root) / f
        if not path.exists():
            missing.append(str(path))
    
    if missing:
        print("以下文件/目录缺失：")
        for m in missing:
            print(f"  - {m}")
    else:
        print("数据集结构完整！")






class Veri776Dataset(Dataset):   #功能：标准VeRi-776数据集加载器  预处理：包含图像尺寸调整、标准化等transform操作
    def __init__(self, mode="train", transform=None, root_dir=None, cache_images=False, id_to_class=None):
        super().__init__()

        start_time = time.time()


     
        # 设置数据集根目录
        if root_dir:
            self.root_dir = root_dir
        else:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            self.root_dir = os.path.abspath(
                os.path.join(current_dir, r"E:\codes\project\datas\VeRi-776")
            )
       

        # 验证路径有效性
        if not os.path.exists(self.root_dir):
            raise FileNotFoundError(f"数据集路径不存在: {self.root_dir}")

        self.mode = mode
        self.transform = transform
        self.cache_images = cache_images  # 保留参数，但实际不使用缓存
        self.cached_images = None  # 不初始化缓存
        # 初始化存储容器
        self.image_paths = []
        self.vehicle_ids = []
        self.camera_ids = []  # 摄像头ID列表存储
        
        
        # 模式验证
        valid_modes = ["train", "test", "query"]
        if mode not in valid_modes:
            raise ValueError(f"Invalid mode {mode}, expected {valid_modes}")
        
        # 文件映射
        name_file_map = {
            "train": "name_train.txt",
            "test": "name_test.txt", 
            "query": "name_query.txt"
        }
        image_dir_map = {
            "train": "image_train",
            "test": "image_test",
            "query": "image_query"
        }
        
        # 读取文件名列表
        name_file = os.path.join(self.root_dir, name_file_map[mode])
        if not os.path.exists(name_file):
            raise FileNotFoundError(f"{name_file} 不存在!")
        with open(name_file, "r") as f:
            file_names = [line.strip() for line in f.readlines()]
        print(f"共加载 {len(file_names)} 个文件名，耗时 {time.time()-start_time:.2f}s")
        

                 
                    

        # 构建数据列表（带进度条）
        image_dir = os.path.join(self.root_dir, image_dir_map[mode])
        self.data = []  # 存储 (路径, 车辆ID, 相机ID) 三元组

        for file_name in tqdm(file_names, desc=f"处理 {mode} 数据"):
            # 解析车辆ID
            try:
                vehicle_id = int(file_name.split("_")[0])
            except:
                print(f"无法解析文件名: {file_name}")
                continue



            # +++ 摄像头ID解析部分 +++
            try:
                cam_part = file_name.split("_")[1]    # 例如提取c002
                cam_id = int(cam_part[1:])            # 转换为整数2
            except Exception as e:
                print(f"无法解析摄像头ID: {file_name}, 错误: {str(e)}")
                cam_id = 0
            
            



            # 检查文件存在性
            img_path = os.path.join(image_dir, file_name)
            if not os.path.exists(img_path):
                print(f"文件缺失: {img_path}")
                continue
                
            # 存储三元组：(路径, 车辆ID, 相机ID)
            self.data.append((img_path, vehicle_id, cam_id))
            self.vehicle_ids.append(vehicle_id)
            self.camera_ids.append(cam_id)

            # 新增：用于存储未知车辆ID的集合
        self.unknown_vehicle_ids = set()


         # ID映射（连续化处理）
        if id_to_class is not None:
            self.id_to_class = id_to_class
            # 生成unique_ids（即使使用全局映射）
            self.unique_ids = sorted(set(self.vehicle_ids))
             # 收集未知车辆ID
            for vid in self.vehicle_ids:
                if vid not in self.id_to_class:
                    self.unknown_vehicle_ids.add(vid)
        else:
            # 仅训练集生成全局映射
            self.unique_ids = sorted(set(self.vehicle_ids))
            self.id_to_class = {v: k for k, v in enumerate(self.unique_ids)}


         # 检查ID连续性
        if mode == "train":
            class_ids = [self.id_to_class[vid] for vid in self.vehicle_ids]
            min_class = min(class_ids)
            max_class = max(class_ids)
            if min_class != 0 or max_class != len(self.unique_ids) - 1:
                print(f"警告: 训练集类别ID不连续，范围: [{min_class}, {max_class}]")


                # 在初始化最后统一输出未知车辆统计信息
        if self.unknown_vehicle_ids and mode != "train":
            print(f"警告: {mode}集发现 {len(self.unknown_vehicle_ids)} 个未知车辆ID，ID范围: {min(self.unknown_vehicle_ids)}-{max(self.unknown_vehicle_ids)}")



     
   
        
        


        



    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        img_path, vehicle_id, cam_id = self.data[idx]
        img = Image.open(img_path).convert("RGB")  # 动态加载图像



        if self.transform:
            img = self.transform(img)
        




        return {
            "image": img,
            "vehicle_id": vehicle_id,
            "class_id": self.id_to_class[vehicle_id],
            "camera_id": cam_id,
            "path": img_path
            
        }



if __name__ == "__main__":


    verify_dataset()
    


    # 极简模式测试
    train_transform = transforms.Compose([
        transforms.Resize(256),  # 先调整为稍大的尺寸
    transforms.CenterCrop(224),  # 从中心裁剪出的区域
        transforms.ToTensor()
    ])
     # 加载训练集，生成全局ID映射
    train_dataset = Veri776Dataset(mode="train", transform=train_transform)
    global_id_to_class = train_dataset.id_to_class  # 保存训练集的全局映射

     # 加载测试集和查询集，复用训练集的映射
    datasets = {}
    # 测试数据集加载
    for mode in ["train", "query", "test"]:
        print(f"\n=== 测试 {mode} 数据集 ===")
        dataset = Veri776Dataset(mode=mode, 
                                 transform=train_transform ,
                                 id_to_class=global_id_to_class  # 传入全局映射
                                 
                                 )
        datasets[mode] = dataset
        # 检查相机ID分布
        camera_counts = Counter(dataset.camera_ids)
        print(f"摄像头ID分布: {camera_counts}")



  

     

    print("=== 测试开始 ===")
    try:
        
        loader = torch.utils.data.DataLoader(
            datasets["train"],  # 指定使用训练集
            batch_size=16,
            num_workers=4,
            shuffle=True
        )
        batch = next(iter(loader))
        print("数据加载测试成功！")
        print("图像尺寸:", batch["image"].shape)
        # 验证ID连续性
        print("\n=== ID连续性验证 ===")
        print("训练集车辆ID范围:", min(train_dataset.vehicle_ids), max(train_dataset.vehicle_ids))
        print("训练集类别ID范围:", 0, len(train_dataset.unique_ids) - 1)  # 应连续0~N-1
        print("训练集相机ID范围:", min(train_dataset.camera_ids), max(train_dataset.camera_ids))

        
        
    except Exception as e:
        print(f"错误发生: {str(e)}")






#警告: query集发现 200 个未知车辆ID，ID范围: 2-776
#警告: test集发现 200 个未知车辆ID，ID范围: 2-776
#VeRi-776 数据集的训练集、查询集、测试集是按摄像头划分的（而非按车辆 ID），因此：
#训练集包含部分摄像头的车辆图像；
#查询集和测试集包含其他摄像头的车辆图像，其中许多车辆未在训练集中出现。
#代码强制要求测试 / 查询集的车辆 ID 必须存在于训练集的映射中，不存在时输出警告。
# 这些警告是 Veri776 数据集开放式测试的正常现象，不影响模型训练和评估结果

#开放式测试：
#测试集包含训练集未出现的车辆，更符合实际场景（如监控场景新增车辆），Veri776 采用此设置。
#模型需要学习车辆特征的泛化能力，而非死记硬背训练集车辆。